package com.example.leetcode.design;

import java.util.*;

/**
 * 超时
 */
class MKAverage0 {

    private final int m;
    private final int k;
    private final List<Integer> container;
    int sum;

    public MKAverage0(int m, int k) {
        this.m = m;
        this.k = k;
        this.container = new ArrayList<>();
        this.sum = 0;
    }

    public void addElement(int num) {
        if (container.size() < m - 1) {
            container.add(num);
        } else {
            // 去掉最早的元素
            if (container.size() == m) {
                container.remove(0);
            }
            // 加入新元素
            container.add(num);
            List<Integer> list = new ArrayList<>(container);
            sum = 0;
            // 升序排序
            Collections.sort(list);
            for (int i = k; i <= m - 1 - k; i++) {
                sum += list.get(i);
            }
        }
    }

    public int calculateMKAverage() {
        if (container.size() < m) {
            return -1;
        }
        double avg = (sum * 1.0) / (m - 2 * k);
        return (int) Math.floor(avg);
    }
}

/**
 * Your MKAverage object will be instantiated and called as such:
 * MKAverage obj = new MKAverage(m, k);
 * obj.addElement(num);
 * int param_2 = obj.calculateMKAverage();
 */

class MKAverage {
    private final int m;
    private final int k;
    private final Queue<Integer> queue;
    private final TreeMap<Integer, Integer> map1;
    private final TreeMap<Integer, Integer> map2;
    private final TreeMap<Integer, Integer> map3;
    private int size1;
    private int size2;
    private long sum;

    public MKAverage(int m, int k) {
        this.m = m;
        this.k = k;
        this.queue = new ArrayDeque<>();
        this.map1 = new TreeMap<>();
        this.map2 = new TreeMap<>();
        this.map3 = new TreeMap<>();
        this.size1 = 0;
        this.size2 = 0;
        this.sum = 0;
    }

    public void addElement(int num) {
        queue.offer(num);
        if (queue.size() <= m) {
            map2.put(num, map2.getOrDefault(num, 0) + 1);
            sum += num;
            if (queue.size() == m) {
                while (size1 < k) {
                    int firstNum = map2.firstKey();
                    map1.put(firstNum, map1.getOrDefault(firstNum, 0) + 1);
                    size1++;
                    sum -= firstNum;
                    map2.put(firstNum, map2.get(firstNum) - 1);
                    if (map2.get(firstNum) == 0) {
                        map2.remove(firstNum);
                    }
                }
                while (size2 < k) {
                    int lastNum = map2.lastKey();
                    map3.put(lastNum, map3.getOrDefault(lastNum, 0) + 1);
                    size2++;
                    sum -= lastNum;
                    map2.put(lastNum, map2.get(lastNum) - 1);
                    if (map2.get(lastNum) == 0) {
                        map2.remove(lastNum);
                    }
                }
            }
            return;
        }

        if (num < map1.lastKey()) {
            map1.put(num, map1.getOrDefault(num, 0) + 1);
            int lastNum = map1.lastKey();
            map2.put(lastNum, map2.getOrDefault(lastNum, 0) + 1);
            sum += lastNum;
            map1.put(lastNum, map1.get(lastNum) - 1);
            if (map1.get(lastNum) == 0) {
                map1.remove(lastNum);
            }
        } else if (num > map3.firstKey()) {
            map3.put(num, map3.getOrDefault(num, 0) + 1);
            int firstNum = map3.firstKey();
            map2.put(firstNum, map2.getOrDefault(firstNum, 0) + 1);
            sum += firstNum;
            map3.put(firstNum, map3.get(firstNum) - 1);
            if (map3.get(firstNum) == 0) {
                map3.remove(firstNum);
            }
        } else {
            map2.put(num, map2.getOrDefault(num, 0) + 1);
            sum += num;
        }

        Integer x = queue.poll();
        if (map1.containsKey(x)) {
            map1.put(x, map1.get(x) - 1);
            if (map1.get(x) == 0) {
                map1.remove(x);
            }
            int firstNum = map2.firstKey();
            map1.put(firstNum, map1.getOrDefault(firstNum, 0) + 1);
            sum -= firstNum;
            map2.put(firstNum, map2.get(firstNum) - 1);
            if (map2.get(firstNum) == 0) {
                map2.remove(firstNum);
            }
        } else if (map3.containsKey(x)) {
            map3.put(x, map3.get(x) - 1);
            if (map3.get(x) == 0) {
                map3.remove(x);
            }
            int lastNum = map2.lastKey();
            map3.put(lastNum, map3.getOrDefault(lastNum, 0) + 1);
            sum -= lastNum;
            map2.put(lastNum, map2.get(lastNum) - 1);
            if (map2.get(lastNum) == 0) {
                map2.remove(lastNum);
            }
        } else {
            map2.put(x, map2.get(x) - 1);
            if (map2.get(x) == 0) {
                map2.remove(x);
            }
            sum -= x;
        }
    }

    public int calculateMKAverage() {
        if (queue.size() < m) {
            return -1;
        }
        return (int) (sum / (m - 2 * k));
    }
}

public class leetcode1825 {
    public static void main(String[] args) {
        MKAverage obj = new MKAverage(3, 1);
        obj.addElement(3);        // 当前元素为 [3]
        obj.addElement(1);        // 当前元素为 [3,1]
        System.out.println(obj.calculateMKAverage()); // 返回 -1 ，因为 m = 3 ，但数据流中只有 2 个元素
        obj.addElement(10);       // 当前元素为 [3,1,10]
        System.out.println(obj.calculateMKAverage()); // 最后 3 个元素为 [3,1,10]
        // 删除最小以及最大的 1 个元素后，容器为 [3]
        // [3] 的平均值等于 3/1 = 3 ，故返回 3
        obj.addElement(5);        // 当前元素为 [3,1,10,5]
        obj.addElement(5);        // 当前元素为 [3,1,10,5,5]
        obj.addElement(5);        // 当前元素为 [3,1,10,5,5,5]
        System.out.println(obj.calculateMKAverage()); // 最后 3 个元素为 [5,5,5]
        // 删除最小以及最大的 1 个元素后，容器为 [5]
        // [5] 的平均值等于 5/1 = 5 ，故返回 5
    }
}
